module CompilerLLVM where

import Control.Monad
import Control.Monad.State
import Data.List (intercalate)
import qualified Data.Map as Map
import Data.Map (Map, empty, insert, fromList)
import AbsJavalette

data Env = Env {
	var :: [Map Ident String],
	regC :: Int,
	labelC :: Int,
	tempCode :: String,
	code :: [String],
	strings :: String
}

type EnvM = State Env
--A label must be followed by an instruction, use unreachable if necessary


emptyState :: Env
emptyState = Env [] 0 0 "" [] []


{-
String literals needs to be stored globally, this function adds the global declaration to string 
and gives back the string name like @string01
c"Hello World"
-}
addString :: String -> EnvM (String)
addString sValue = do
	Env v rc lc tc c s <- get
	let stringName = "@string" ++ show rc
	let stringLength = show ((length sValue) + 1)
	let stringDeclaration =  stringName ++ " = internal constant [" ++ stringLength ++ " x i8] c\"" ++ sValue ++ "\\00\"\n"
	put (Env v (rc+1) lc tc c (s++stringDeclaration))
	return stringName

getStrings :: EnvM (String)
getStrings = do
	Env _ _ _ _ _ s <- get
	return s

getNextRegNo :: EnvM (String)
getNextRegNo = do
	Env v rc lc tc c s <- get
	put (Env v (rc+1) lc tc c s)
	return (show rc)

addVar :: Ident -> String -> EnvM ()
addVar identifier reg= do
	Env (v:vs) rc lc tc c s <- get
	let v' = (insert identifier reg v)
	put (Env (v':vs) rc lc tc c s)

getLabel :: EnvM (String)
getLabel = do
	Env v rc lc tc c s<- get
	put (Env v rc (lc+1) tc c s)
	return ("lbl" ++ show lc)

lookupVar :: Ident -> EnvM (String)
lookupVar identifier = do
	Env v _ _ _ _ _ <- get
	return (findVarInStack v identifier)

findVarInStack :: [Map Ident String] -> Ident -> String
findVarInStack [] (Ident identifier) = "Fuck it " ++ identifier ++ " does not exist"
findVarInStack (t:rest) identifier = case (Map.lookup identifier t) of
	Just i -> i
	Nothing -> findVarInStack rest identifier

getCode :: EnvM ((String, String))
getCode = do
	Env v rc lc tc c s <- get
	return ((concatenation c), s)

newRegister :: EnvM (String)
newRegister = do
	Env v rc lc tc c s<- get
	put (Env v (rc+1) lc tc c s)
	return ("%t" ++ show rc)

newBlock :: EnvM ()
newBlock = do
	Env v rc lc tc c s<- get
	put (Env (empty:v) rc lc tc c s)

exitBlock :: EnvM ()
exitBlock = do
	Env (v:vs) rc lc tc c s<- get
	put (Env vs rc lc tc c s)

getType :: Expr -> String
getType (ETyped Int _) = "i32"
getType (ETyped Doub _) = "double"
getType (ETyped Void _) = ""
getType (ETyped Bool _) = "i1"


--["%t1", "%t2", "5",..]
--"i32 %t1, double %t2, i32 5"
compileArgs :: [Expr] -> [String] -> EnvM ([String])
compileArgs [] l = return l
compileArgs (e:es) s = do
	er <- compileExpr e
	compileArgs es (s++[((getType e) ++ " " ++ er)])


compileItems :: Type -> [Item] -> EnvM ()
compileItems t [] = return ()
compileItems t (i:is) = do
	case i of
		NoInit identifier@(Ident name) -> do
			case t of
				Int -> do
					n <- getNextRegNo
					let vReg = "%" ++ name ++ n
					addVar identifier vReg
					addTemporaryCode (vReg ++ " = alloca i32\n")
					addTemporaryCode ("store i32 0 , i32* " ++ vReg ++ "\n")
				Doub -> do
					n <- getNextRegNo
					let vReg = "%" ++ name ++ n
					addVar identifier vReg
					addTemporaryCode (vReg ++ " = alloca double\n")
					addTemporaryCode ("store double 0.0 , double* " ++ vReg ++ "\n")
				Bool -> do
					n <- getNextRegNo
					let vReg = "%" ++ name ++ n
					addVar identifier vReg
					addTemporaryCode (vReg ++ " = alloca i1\n")
					addTemporaryCode ("store i1 0 , i1* " ++ vReg ++ "\n")
		Init identifier@(Ident name) e -> do
			case t of
				Int -> do
					n <- getNextRegNo
					let vReg = "%" ++ name ++ n
					v <- compileExpr e
					addVar identifier vReg
					addTemporaryCode (vReg ++ " = alloca i32\n")
					addTemporaryCode ("store i32 " ++ v ++ ", i32* " ++ vReg ++ "\n")
				Doub -> do
					n <- getNextRegNo
					let vReg = "%" ++ name ++ n
					v <- compileExpr e
					addVar identifier vReg
					addTemporaryCode (vReg ++ " = alloca double\n")
					addTemporaryCode ("store double " ++ v ++ ", double* " ++ vReg ++ "\n")
				Bool -> do
					n <- getNextRegNo
					let vReg = "%" ++ name ++ n
					v <- compileExpr e
					addVar identifier vReg
					addTemporaryCode (vReg ++ " = alloca i1\n")	
					addTemporaryCode ("store i1 " ++ v ++ ", i1* " ++ vReg ++ "\n")
	compileItems t is

addTemporaryCode :: String -> EnvM ()
addTemporaryCode newCode = do
	Env v rc lc tc c s<- get
	put (Env v rc lc (tc++newCode) c s)

compileProgram :: Program -> String
compileProgram (Program td) = do
	let (code, string) = evalState (compileFunctions td >> getCode) emptyState
	fileDeclaration code string

compileFunctionArguments :: [Arg] -> [String] -> EnvM ([String])
compileFunctionArguments [] s = return (s)
compileFunctionArguments (a@(Arg t identifier@(Ident name)):as) s = do
	reg <- newRegister
	n <- getNextRegNo
	let tName = case t of
		Int -> "i32"
		Doub -> "double"
		Bool -> "i1"
	let s' = tName ++ reg
	let vReg = "%" ++ name ++ n
	addVar identifier vReg
	case t of
		Int -> do
			addTemporaryCode (vReg ++ " = alloca i32\n")
			addTemporaryCode ("store i32 " ++ reg ++ ", i32* " ++ vReg ++ "\n")
		Doub -> do
			addTemporaryCode (vReg ++ " = alloca double\n")
			addTemporaryCode ("store double " ++ reg ++ ", double* " ++ vReg ++ "\n")
		Bool -> do
			addTemporaryCode (vReg ++ " = alloca i1\n")
			addTemporaryCode ("store i1 " ++ reg ++ ", i1* " ++ vReg ++ "\n")
	compileFunctionArguments as (s ++ [s'])

compileFunctions :: [TopDef] -> EnvM ()
compileFunctions [] = return ()
compileFunctions (f:fs) = do
	compileFunction f
	compileFunctions fs
	return ()


compileFunction :: TopDef -> EnvM ()
compileFunction f@(FnDef t (Ident i) args (Block b)) = do
	newBlock
	s <- compileFunctionArguments args []
	let s' = intercalate "," s
	compileStms b
	doFinalCode f s'
	return ()

doFinalCode:: TopDef -> String -> EnvM ()
doFinalCode (FnDef typ (Ident identifier) args (Block b)) s'= do
	hackRet typ
	Env v rc lc tc c s<- get
	let t = case typ of
		Void -> "void"
		Int -> "i32"
		Doub -> "double"
		Bool -> "i1"
	let newFinal = c ++ ["define " ++ t ++ "@" ++ identifier ++ "(" ++ s' ++ "){\nentry:\n" ++ tc ++ "}\n\n"]
	put (Env [] rc lc "" newFinal s)

compileStms :: [Stmt] -> EnvM ()
compileStms [] = return ()
compileStms (s:ss) = do
	compileStm s
	compileStms ss
	return ()

compileStm :: Stmt -> EnvM ()
compileStm stmt = case stmt of
	Empty -> return ()
	BStmt (Block stmts) -> do
		newBlock
		compileStms stmts
		exitBlock
	Ret e -> do
		ev <- compileExpr e
		addTemporaryCode ("ret " ++ (getType e) ++ " " ++ ev ++ "\n")
	VRet -> do
		addTemporaryCode ("ret void\n")
	SExp e -> do
		compileExpr e
		return ()
	Decl t items -> do
		compileItems t items
		return ()
	Ass identifier e -> do
		sl <- lookupVar identifier
		res <- compileExpr e
		addTemporaryCode ("store " ++ (getType e) ++ " " ++ res ++ ", " ++ (getType e) ++ "* " ++ sl ++ "\n")
	Incr identifier -> do
		sl <- lookupVar identifier
		reg0 <- newRegister
		reg1 <- newRegister
		addTemporaryCode (reg0 ++ " = load i32* " ++ sl ++ "\n")
		addTemporaryCode (reg1 ++ " = add i32 1, " ++ reg0 ++ "\n")
		addTemporaryCode ("store i32 " ++ reg1 ++ ", i32* " ++ sl ++ "\n")
	Decr identifier -> do
		sl <- lookupVar identifier
		reg0 <- newRegister
		reg1 <- newRegister
		addTemporaryCode (reg0 ++ " = load i32* " ++ sl ++ "\n")
		addTemporaryCode (reg1 ++ " = sub i32 " ++ reg0 ++ ", 1 \n")
		addTemporaryCode ("store i32 " ++ reg1 ++ ", i32* " ++ sl ++ "\n")
	Cond e@(ETyped t innerE) s -> do
		case innerE of
			ELitTrue -> do
				compileStm s
				return ()
			ELitFalse -> do
				return ()
			_ -> do
				falseLabel <- getLabel
				trueLabel <- getLabel
				eReg <- compileExpr e
				addTemporaryCode ("br i1 " ++ eReg ++ ", label %" ++ trueLabel ++ ", label %" ++ falseLabel ++ "\n")
				addTemporaryCode (trueLabel ++ ":\n")
				compileStm s
				addTemporaryCode ("br label %" ++ falseLabel ++ "\n")
				addTemporaryCode (falseLabel ++ ":\n")
	CondElse e@(ETyped t innerE) s1 s2 -> do
		case innerE of
			ELitTrue -> do
				compileStm s1
				return ()
			ELitFalse -> do
				compileStm s2
				return ()
			_ -> do
				falseLabel <- getLabel
				trueLabel <- getLabel
				endLabel <- getLabel
				reg <- compileExpr e
				addTemporaryCode ("br i1 " ++ reg ++ ", label %" ++ trueLabel ++ ", label %" ++ falseLabel ++ "\n")
				addTemporaryCode (trueLabel ++ ":\n")
				compileStm s1
				addTemporaryCode ("br label %" ++ endLabel ++ "\n")
				addTemporaryCode (falseLabel ++ ":\n")
				compileStm s2
				addTemporaryCode ("br label %" ++ endLabel ++ "\n")
				addTemporaryCode (endLabel ++ ":\n")
	While e@(ETyped t innerE) s -> do
		case innerE of 
			ELitFalse -> do
				return ()
			ELitTrue -> do
				beginLabel <- getLabel
				addTemporaryCode ("br label %" ++ beginLabel ++ "\n")
				addTemporaryCode (beginLabel ++ ":\n")
				compileStm s
				addTemporaryCode ("br label %" ++ beginLabel ++ "\n")
				return ();
			_ -> do
				beginLabel <- getLabel
				beforeStmt <- getLabel
				endLabel <- getLabel
				addTemporaryCode ("br label %" ++ beginLabel ++ "\n")
				addTemporaryCode (beginLabel ++ ":\n")
				reg <- compileExpr e
				addTemporaryCode ("br i1 " ++ reg ++ ", label %" ++ beforeStmt ++ ", label %" ++ endLabel ++ "\n")
				addTemporaryCode (beforeStmt ++ ":\n")
				compileStm s
				addTemporaryCode ("br label %" ++ beginLabel ++ "\n")
				addTemporaryCode (endLabel ++ ":\n")
				return ()





compileExpr :: Expr -> EnvM (String)
compileExpr expr = case expr of
	ETyped t (ELitInt v) -> do
		return (show v)
	ETyped t (ELitDoub v) -> do
		return (show v)
	ETyped t (ELitTrue) -> do
		return (show 1)
	ETyped t (ELitFalse) -> do
		return (show 0)
	ETyped _ (EApp (Ident "printString") args) -> do
		let outString =  typedEStringToString (head args)
		sLoc <- addString outString
		let stringLength = show (length outString + 1)
		sReg <- newRegister
		addTemporaryCode (sReg ++ " = bitcast [" ++ stringLength ++ " x i8]* " ++ sLoc ++ " to i8*\n")
		addTemporaryCode ("call void @printString(i8* " ++ sReg ++ ")\n")
		return ("")
	ETyped t (EApp (Ident identifier) args) -> do
		a <- compileArgs args []
		let a' = intercalate "," a
		case t of
			Void -> do
				addTemporaryCode ("call void @" ++ identifier ++ "("++ a' ++ ")\n")
				return ("")
			Int -> do
				reg <- newRegister
				addTemporaryCode (reg ++ " = call i32 @" ++ identifier ++ "("++ a' ++ ")\n")
				return reg
			Doub -> do
				reg <- newRegister
				addTemporaryCode (reg ++ " = call double @" ++ identifier ++ "("++ a' ++ ")\n")
				return reg
			Bool -> do
				reg <- newRegister
				addTemporaryCode (reg ++ " = call i1 @" ++ identifier ++ "("++ a' ++ ")\n")
				return reg
	ETyped t (EAdd a op b) -> case op of
		Plus -> do
			reg <- newRegister
			al <- compileExpr a
			bl <- compileExpr b
			case t of
				Int -> addTemporaryCode (reg ++ " = add " ++ (getType a) ++ " " ++ al ++ ", " ++ bl ++ "\n")
				Doub -> addTemporaryCode (reg ++ " = fadd " ++ (getType a) ++ " " ++ al ++ ", " ++ bl ++ "\n")
			return (reg)
		Minus -> do
			reg <- newRegister
			al <- compileExpr a
			bl <- compileExpr b
			case t of
				Int -> addTemporaryCode (reg ++ " = sub " ++ (getType a) ++ " " ++ al ++ ", " ++ bl ++ "\n")
				Doub -> addTemporaryCode (reg ++ " = fsub " ++ (getType a) ++ " " ++ al ++ ", " ++ bl ++ "\n")
			return (reg)
	ETyped t (EMul a op b) -> case op of
		Times -> do
			reg <- newRegister
			al <- compileExpr a
			bl <- compileExpr b
			case t of
				Int -> addTemporaryCode (reg ++ " = mul " ++ (getType a) ++ " " ++ al ++ ", " ++ bl ++ "\n")
				Doub -> addTemporaryCode (reg ++ " = fmul " ++ (getType a) ++ " " ++ al ++ ", " ++ bl ++ "\n")
			return reg
		Div -> do
			reg <- newRegister
			al <- compileExpr a
			bl <- compileExpr b
			case t of
				Int -> addTemporaryCode (reg ++ " = sdiv " ++ (getType a) ++ " " ++ al ++ ", " ++ bl ++ "\n")
				Doub -> addTemporaryCode (reg ++ " = fdiv " ++ (getType a) ++ " " ++ al ++ ", " ++ bl ++ "\n")
			return reg
		Mod -> do
			reg <- newRegister
			al <- compileExpr a
			bl <- compileExpr b
			case t of
				Int -> addTemporaryCode (reg ++ " = srem " ++ (getType a) ++ " " ++ al ++ ", " ++ bl ++ "\n")
			return reg
	ETyped t (Neg e) -> do
		reg <- newRegister
		el <- compileExpr e
		case t of
			Int -> do
				addTemporaryCode (reg ++ " = sub i32 0, " ++ el ++ "\n")
			Doub -> do
				addTemporaryCode (reg ++ " = fsub double 0.0, " ++ el ++ "\n")
		return reg
	ETyped t (Not e) -> do
		reg <- newRegister
		el <- compileExpr e
		addTemporaryCode (reg ++ " = xor i1 1, " ++ el ++ "\n")
		return reg
	ETyped t (EVar i) -> do
		sl <- lookupVar i
		reg <- newRegister
		case t of
			Int -> addTemporaryCode (reg ++ " = load i32* " ++ sl ++ "\n")
			Doub -> addTemporaryCode (reg ++ " = load double* " ++ sl ++ "\n")
			Bool -> addTemporaryCode (reg ++ " = load i1* " ++ sl ++ "\n")
		return reg
	ETyped t (ERel e1 op e2) -> do
		e1R <- compileExpr e1
		e2R <- compileExpr e2
		r <- newRegister
		case typeOf e1 of
			Int -> do
				addTemporaryCode (r ++ " = icmp " ++ (compTypeInt op) ++ " i32 " ++ e1R ++ ", " ++ e2R ++ "\n")
				return r
			Bool -> do
				addTemporaryCode (r ++ " = icmp " ++ (compTypeInt op) ++ " i1 " ++ e1R ++ ", " ++ e2R ++ "\n")
				return r
			Doub -> do
				addTemporaryCode (r ++ " = fcmp " ++ (compTypeDoub op) ++ " double " ++ e1R ++ ", " ++ e2R ++ "\n")
				return r
	ETyped t (EAnd e1@(ETyped t1 innerE1) e2@(ETyped t2 innerE2)) -> do
		reg <- newRegister
		if innerE1 == ELitFalse || innerE2 == ELitFalse
			then
				do
					addTemporaryCode (reg ++ " = i1 0\n")
					return reg
			else
				do
					tmpVar <- newRegister --holds the return value
					endLabel <- getLabel
					beginLabel <- getLabel
					continueLabel1 <- getLabel
					continueLabel2 <- getLabel
					addTemporaryCode (tmpVar ++ " = alloca i1\n")
					addTemporaryCode ("store i1 0, i1* " ++ tmpVar ++ "\n")
					----
					reg1 <- compileExpr e1
					addTemporaryCode ("br i1 " ++ reg1 ++ ", label %" ++ continueLabel1 ++ ", label %" ++ endLabel ++ "\n")
					addTemporaryCode (continueLabel1 ++ ":\n")
					reg2 <- compileExpr e2
					addTemporaryCode ("br i1 " ++ reg2 ++ ", label %" ++ continueLabel2 ++ ", label %" ++ endLabel ++ "\n")
					addTemporaryCode (continueLabel2 ++ ":\n")
					addTemporaryCode ("store i1 1, i1* " ++ tmpVar ++ "\n")
					addTemporaryCode ("br label %" ++ endLabel ++ "\n")
					addTemporaryCode (endLabel ++ ":\n")
					addTemporaryCode (reg ++ " = load i1* " ++ tmpVar ++ "\n")
					return reg
	ETyped t (EOr e1@(ETyped t1 innerE1) e2@(ETyped t2 innerE2)) -> do
		reg <- newRegister
		if innerE1 == ELitTrue  ||  innerE2 == ELitTrue
			then 
				do
					addTemporaryCode (reg ++ " = i1 1\n")
					return reg
			else
				do
					tmpVar <- newRegister --holds the return value
					endLabel <- getLabel
					beginLabel <- getLabel
					continueLabel1 <- getLabel
					continueLabel2 <- getLabel
					addTemporaryCode (tmpVar ++ " = alloca i1\n")
					addTemporaryCode ("store i1 1, i1* " ++ tmpVar ++ "\n")
					----
					reg1 <- compileExpr e1
					addTemporaryCode ("br i1 " ++ reg1 ++ ", label %" ++ endLabel ++ ", label %" ++ continueLabel1 ++ "\n")
					addTemporaryCode (continueLabel1 ++ ":\n")
					reg2 <- compileExpr e2
					addTemporaryCode ("br i1 " ++ reg2 ++ ", label %" ++ endLabel ++ ", label %" ++ continueLabel2 ++ "\n")
					addTemporaryCode (continueLabel2 ++ ":\n")
					addTemporaryCode ("store i1 0, i1* " ++ tmpVar ++ "\n")
					addTemporaryCode ("br label %" ++ endLabel ++ "\n")
					addTemporaryCode (endLabel ++ ":\n")
					addTemporaryCode (reg ++ " = load i1* " ++ tmpVar ++ "\n")
					return reg


fileDeclaration :: String -> String -> String
fileDeclaration code strings= strings ++ "declare void @printInt(i32)\n" ++
	"declare void @printDouble(double)\n" ++
	"declare void @printString(i8*)\n" ++
	"declare i32 @readInt()\n" ++
	"declare double @readDouble()\n\n" ++ code


concatenation :: [String] -> String
concatenation [] = ""
concatenation (s:ss) = s ++ concatenation ss

typeOf :: Expr -> Type
typeOf (ETyped t _) = t

compTypeInt :: RelOp -> String
compTypeInt op = case op of
	LTH -> "sle"
	LE -> "slt"
	GTH -> "sgt"
	GE -> "sge"
	EQU -> "eq"
	NE -> "ne"

compTypeDoub :: RelOp -> String
compTypeDoub op = case op of
	LTH -> "ole"
	LE -> "olt"
	GTH -> "ogt"
	GE -> "oge"
	EQU -> "oeq"
	NE -> "one"

--Hack to get the string needed for application of the printString function
typedEStringToString :: Expr -> String
typedEStringToString (ETyped String (EString s)) = s

hackRet :: Type -> EnvM ()
hackRet t = case t of
	Void -> addTemporaryCode "ret void\n"
	Int -> addTemporaryCode "ret i32 0\n"
	Doub -> addTemporaryCode "ret double 0.0\n"
	Bool -> addTemporaryCode "ret i1 0\n"



	